"""DHF模型参数配置模块

包含产流参数、汇流参数和蒸发参数的定义和管理。
"""

import numpy as np
from typing import Dict, Any, Optional
import json
from pathlib import Path
from loguru import logger


class DHFParameters:
    """DHF模型参数类"""
    
    def __init__(self, initial_flow: float = 50.0):
        """
        初始化DHF模型参数
        
        Args:
            initial_flow: 降雨开始时的起始流量 (m³/s)
        """
        self.initial_flow = initial_flow
        self._setup_runoff_parameters()
        self._setup_routing_parameters()
        self._setup_evaporation_parameters()
    
    def _setup_runoff_parameters(self):
        """设置产流参数"""
        if self.initial_flow >= 100.0:
            # 表10-1: 降雨起始流量 ≥100 m³/s 的参数
            self.runoff_params = {
                'g': 0.05,      # 不透水面积占全流域面积之比值
                'A': 4.129,     # 表层蓄水量曲线形状参数
                'B': 1.181,     # 下渗抛物线分配曲线形状参数
                'K': 0.28,      # 下层下渗曲线曲率
                'S0': 45,       # 表层蓄水容量 (mm)
                'U0': 100,      # 下层蓄水容量 (mm)
                'D0': 130       # 下层蓄水容量与地下水库蓄水容量之和 (mm)
            }
        else:
            # 表10-2: 降雨起始流量 <100 m³/s 的参数
            self.runoff_params = {
                'g': 0.05,
                'A': 4.129,
                'B': 1.371,
                'K': 0.469,
                'S0': 80,
                'U0': 100,
                'D0': 140
            }
        
        # 计算V0 (地下水库蓄水容量)
        self.runoff_params['V0'] = self.runoff_params['D0'] - self.runoff_params['U0']
        
        # 湿润地区地下径流系数
        self.runoff_params['K_W'] = 1.0
    
    def _setup_routing_parameters(self):
        """设置汇流参数"""
        # 表10-3: 大伙房模型汇流参数
        self.routing_params = {
            'Ka': 0.8,      # 前期净雨影响程度衰减系数
            'L': 120,       # 流域最大河长 (km)
            'N': 2.5,       # 地下径流与地面壤中流汇流曲线底宽的比例系数
            'B0': 0.8,      # 特征河长计算参数
            'K0': 0.4,      # 拟合线的斜率
            'Cc1': 0.3,     # 汇流曲线形状参数1
            'Dd1': 0.7      # 汇流曲线形状参数1
        }
        
        # 根据降雨中心位置的不同参数
        self.rainfall_center_params = {
            '上游': {'Cc': 0.8, 'Dd': 2.8, 'Coe': 0.19},
            '中游': {'Cc': 1.0, 'Dd': 1.2, 'Coe': 0.13},
            '下游': {'Cc': 0.6, 'Dd': 1.4, 'Coe': 0.31}
        }
        
        # 计算特征河长
        self.routing_params['L_B'] = self.routing_params['L'] / self.routing_params['B0']
    
    def _setup_evaporation_parameters(self):
        """设置蒸发参数"""
        if self.initial_flow >= 100.0:
            # 表10-1续表: 流域蒸发能力 Ed (mm/d)
            self.evaporation_params = {
                6: {  # 6月
                    'low': 3.7,    # 日雨量 < 5mm
                    'medium': 1.9, # 5mm ≤ 日雨量 < 30mm
                    'high': 1.0    # 日雨量 ≥ 30mm
                },
                7: {  # 7月
                    'low': 4.3,
                    'medium': 1.9,
                    'high': 1.0
                },
                8: {  # 8月
                    'low': 4.3,
                    'medium': 1.9,
                    'high': 1.0
                },
                9: {  # 9月
                    'low': 3.7,
                    'medium': 1.9,
                    'high': 1.0
                }
            }
        else:
            # 表10-2续表: 流域蒸发能力 Ed (mm/d)
            self.evaporation_params = {
                6: {
                    'low': 3.0,
                    'medium': 2.3,
                    'high': 1.4
                },
                7: {
                    'low': 5.1,
                    'medium': 2.3,
                    'high': 1.4
                },
                8: {
                    'low': 5.1,
                    'medium': 2.3,
                    'high': 1.4
                },
                9: {
                    'low': 3.0,
                    'medium': 2.3,
                    'high': 1.4
                }
            }
        
        # 流域蒸发能力与大水体蒸发量的比值
        self.evaporation_params['K_C'] = 1.0
    
    def get_evaporation_capacity(self, month: int, daily_rainfall: float) -> float:
        """
        根据月份和日降雨量获取流域蒸发能力
        
        Args:
            month: 月份 (6-9)
            daily_rainfall: 日降雨量 (mm)
            
        Returns:
            流域蒸发能力 (mm/d)
        """
        if month not in self.evaporation_params:
            raise ValueError(f"不支持的月份: {month}，仅支持6-9月")
        
        month_params = self.evaporation_params[month]
        
        if daily_rainfall < 5.0:
            return month_params['low']
        elif daily_rainfall < 30.0:
            return month_params['medium']
        else:
            return month_params['high']
    
    def get_routing_params_by_center(self, center: str) -> Dict[str, float]:
        """
        根据降雨中心位置获取汇流参数
        
        Args:
            center: 降雨中心位置 ('上游', '中游', '下游')
            
        Returns:
            对应的汇流参数字典
        """
        if center not in self.rainfall_center_params:
            raise ValueError(f"不支持的降雨中心位置: {center}")
        
        return self.rainfall_center_params[center]
    
    def update_parameter(self, param_type: str, param_name: str, value: float):
        """
        更新参数值
        
        Args:
            param_type: 参数类型 ('runoff', 'routing', 'evaporation')
            param_name: 参数名称
            value: 新的参数值
        """
        if param_type == 'runoff':
            if param_name in self.runoff_params:
                self.runoff_params[param_name] = value
            else:
                raise KeyError(f"产流参数中不存在: {param_name}")
        elif param_type == 'routing':
            if param_name in self.routing_params:
                self.routing_params[param_name] = value
            else:
                raise KeyError(f"汇流参数中不存在: {param_name}")
        else:
            raise ValueError(f"不支持的参数类型: {param_type}")
    
    def get_all_parameters(self) -> Dict[str, Any]:
        """
        获取所有参数
        
        Returns:
            包含所有参数的字典
        """
        return {
            'initial_flow': self.initial_flow,
            'runoff_params': self.runoff_params,
            'routing_params': self.routing_params,
            'evaporation_params': self.evaporation_params,
            'rainfall_center_params': self.rainfall_center_params
        }
    
    def __str__(self) -> str:
        """参数信息的字符串表示"""
        return f"DHF模型参数 (起始流量: {self.initial_flow} m³/s)"
    
    def __repr__(self) -> str:
        return self.__str__()
    
    @classmethod
    def from_json(cls, json_path: str, initial_flow: float = None):
        """
        从JSON文件加载参数
        
        Args:
            json_path: JSON参数文件路径
            initial_flow: 起始流量，如果提供则覆盖JSON中的值
            
        Returns:
            DHFParameters实例
        """
        with open(json_path, 'r', encoding='utf-8') as f:
            params_data = json.load(f)
        
        # 创建实例
        if initial_flow is None:
            initial_flow = params_data.get('initial_flow', 50.0)
        
        instance = cls(initial_flow)
        
        # 加载产流参数
        if 'runoff_params' in params_data:
            instance.runoff_params.update(params_data['runoff_params'])
        
        # 加载汇流参数
        if 'routing_params' in params_data:
            instance.routing_params.update(params_data['routing_params'])
        
        # 加载蒸发参数
        if 'evaporation_params' in params_data:
            instance.evaporation_params.update(params_data['evaporation_params'])
        
        # 加载降雨中心参数
        if 'rainfall_center_params' in params_data:
            instance.rainfall_center_params.update(params_data['rainfall_center_params'])
        
        logger.info(f"已从 {json_path} 加载参数")
        return instance
    
    def save_to_json(self, json_path: str):
        """
        保存参数到JSON文件
        
        Args:
            json_path: JSON文件保存路径
        """
        params_data = self.get_all_parameters()
        
        # 确保目录存在
        Path(json_path).parent.mkdir(parents=True, exist_ok=True)
        
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(params_data, f, indent=2, ensure_ascii=False)
        
        logger.info(f"参数已保存到 {json_path}")